#!/usr/bin/env python3
"""
重置向量存储，清理所有集合并重新创建
"""

import sys
import os
import shutil
import logging
from datetime import datetime
from pathlib import Path

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.core.vector_store import get_vector_store
from app.configs.settings import config

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def reset_chroma_database():
    """重置ChromaDB数据库"""
    try:
        logger.info("🗑️ 开始重置ChromaDB数据库...")
        
        # 删除整个chroma_db目录
        chroma_dir = Path("chroma_db")
        if chroma_dir.exists():
            logger.info(f"📁 删除现有数据库目录: {chroma_dir}")
            shutil.rmtree(chroma_dir)
            logger.info("✅ 数据库目录已删除")
        else:
            logger.info("📁 数据库目录不存在，无需删除")
        
        # 重新创建目录
        chroma_dir.mkdir(exist_ok=True)
        logger.info("📁 重新创建数据库目录")
        
        return True
        
    except Exception as e:
        logger.error(f"❌ 重置数据库失败: {e}")
        return False


def create_new_vector_store():
    """创建新的向量存储"""
    try:
        logger.info("🚀 创建新的向量存储...")
        
        # 强制重新初始化向量存储
        vector_store = get_vector_store(force_reinit=True)
        
        logger.info("✅ 新向量存储创建成功")
        return vector_store
        
    except Exception as e:
        logger.error(f"❌ 创建向量存储失败: {e}")
        return None


def import_knowledge_base(vector_store):
    """导入知识库数据"""
    try:
        logger.info("📚 开始导入知识库数据...")
        
        # 导入心理健康知识
        knowledge_dir = Path("data_knowlege")
        if not knowledge_dir.exists():
            logger.warning(f"⚠️ 知识库目录不存在: {knowledge_dir}")
            return False
        
        # 查找所有文本文件
        text_files = list(knowledge_dir.glob("*.txt"))
        md_files = list(knowledge_dir.glob("*.md"))
        all_files = text_files + md_files
        
        logger.info(f"📄 找到 {len(all_files)} 个文件待导入")
        
        if not all_files:
            logger.warning("⚠️ 没有找到知识库文件")
            return False
        
        documents = []
        metadatas = []
        
        for file_path in all_files:
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read().strip()
                
                if content:
                    documents.append(content)
                    metadatas.append({
                        "source": str(file_path),
                        "type": "knowledge_base",
                        "imported_at": datetime.now().isoformat(),
                        "file_type": file_path.suffix
                    })
                    logger.info(f"✅ 读取文件: {file_path.name} ({len(content)} 字符)")
                
            except Exception as e:
                logger.error(f"❌ 读取文件失败 {file_path}: {e}")
        
        if documents:
            logger.info(f"📝 开始向量化并存储 {len(documents)} 个文档...")
            
            # 批量添加文档
            vector_store.add_texts(
                texts=documents,
                metadatas=metadatas
            )
            
            logger.info(f"✅ 成功导入 {len(documents)} 个知识库文档")
            return True
        else:
            logger.warning("⚠️ 没有找到有效的文档内容")
            return False
        
    except Exception as e:
        logger.error(f"❌ 导入知识库失败: {e}")
        return False


def add_sample_conversations(vector_store):
    """添加一些示例对话"""
    try:
        logger.info("💬 添加示例对话...")
        
        sample_conversations = [
            {
                "content": "用户: 我是谁\n助手: 您好！您是一个独特而宝贵的个体。每个人都有自己的身份、经历和价值。如果您想探讨自我认知或身份认同的话题，我很乐意与您一起探讨。",
                "metadata": {
                    "source": "sample_conversation",
                    "type": "conversation",
                    "user_emotion": "neutral",
                    "crisis_level": "low",
                    "timestamp": str(int(datetime.now().timestamp())),
                    "conversation_id": "sample_001",
                    "tools_used": "none"
                }
            },
            {
                "content": "用户: 我最近感到很焦虑，工作压力很大\n助手: 我理解您现在面临的压力。工作焦虑是很常见的情况。建议您可以尝试一些放松技巧，比如深呼吸练习、适当的运动，或者与信任的朋友分享您的感受。如果焦虑持续影响您的日常生活，建议寻求专业心理咨询师的帮助。",
                "metadata": {
                    "source": "sample_conversation",
                    "type": "conversation",
                    "user_emotion": "anxious",
                    "crisis_level": "low",
                    "timestamp": str(int(datetime.now().timestamp())),
                    "conversation_id": "sample_002",
                    "tools_used": "emotion_analysis,professional_advice"
                }
            },
            {
                "content": "用户: 如何应对压力\n助手: 应对压力有很多有效的方法：1. 深呼吸和冥想练习 2. 规律的运动 3. 保持良好的睡眠习惯 4. 与朋友和家人保持联系 5. 学会说'不'，避免过度承诺 6. 寻找兴趣爱好来放松 7. 如果压力过大，寻求专业帮助。记住，适度的压力是正常的，关键是学会管理它。",
                "metadata": {
                    "source": "sample_conversation",
                    "type": "conversation",
                    "user_emotion": "stressed",
                    "crisis_level": "low",
                    "timestamp": str(int(datetime.now().timestamp())),
                    "conversation_id": "sample_003",
                    "tools_used": "professional_advice,knowledge_retrieval"
                }
            }
        ]
        
        documents = [conv["content"] for conv in sample_conversations]
        metadatas = [conv["metadata"] for conv in sample_conversations]
        
        vector_store.add_texts(
            texts=documents,
            metadatas=metadatas
        )
        
        logger.info(f"✅ 成功添加 {len(sample_conversations)} 个示例对话")
        return True
        
    except Exception as e:
        logger.error(f"❌ 添加示例对话失败: {e}")
        return False


def test_vector_store(vector_store):
    """测试向量存储功能"""
    try:
        logger.info("🔍 测试向量存储功能...")
        
        test_queries = ["我是谁", "焦虑症", "如何应对压力"]
        
        for query in test_queries:
            logger.info(f"📝 测试查询: '{query}'")
            results = vector_store.similarity_search(query, k=2)
            
            if results:
                logger.info(f"✅ 找到 {len(results)} 个相关文档")
                for i, doc in enumerate(results):
                    logger.info(f"  {i+1}. {doc.page_content[:50]}...")
            else:
                logger.warning(f"⚠️ 查询 '{query}' 没有找到相关文档")
        
        return True
        
    except Exception as e:
        logger.error(f"❌ 测试向量存储失败: {e}")
        return False


def main():
    """主函数"""
    try:
        logger.info("🎯 开始重置向量存储...")
        
        # 1. 重置数据库
        if not reset_chroma_database():
            logger.error("💥 重置数据库失败")
            sys.exit(1)
        
        # 2. 创建新的向量存储
        vector_store = create_new_vector_store()
        if not vector_store:
            logger.error("💥 创建向量存储失败")
            sys.exit(1)
        
        # 3. 导入知识库
        if not import_knowledge_base(vector_store):
            logger.warning("⚠️ 知识库导入失败，但继续执行")
        
        # 4. 添加示例对话
        if not add_sample_conversations(vector_store):
            logger.warning("⚠️ 示例对话添加失败，但继续执行")
        
        # 5. 测试功能
        if not test_vector_store(vector_store):
            logger.error("💥 向量存储测试失败")
            sys.exit(1)
        
        logger.info("🎉 向量存储重置完成！")
        sys.exit(0)
        
    except Exception as e:
        logger.error(f"❌ 重置过程异常: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()
